# crypto_system.py - A comprehensive cryptographic system implementation

"""
This module implements a custom block cipher with ECB and CBC modes of operation.
The system includes encryption, decryption, key expansion, and utility functions.
"""

import numpy as np
from PIL import Image
import random
import struct
import os
import hashlib
import time
from typing import Tuple, List

# ==================== UTILITY FUNCTIONS ====================

def demo(lst: str, k: int) -> str:
    """
    Circular shift (rotation) of a binary string.
    
    Args:
        lst: Binary string to rotate
        k: Number of positions to rotate
        
    Returns:
        Rotated binary string
        
    Example:
        >>> demo('101100', 2)
        '110010'
    """
    if not lst or k == 0:
        return lst
    k = k % len(lst)
    return lst[k:] + lst[:k]

def xor32(x: str, y: str) -> str:
    """
    32-bit XOR operation between two binary strings.
    
    Args:
        x: First 32-bit binary string
        y: Second 32-bit binary string
        
    Returns:
        XOR result as 32-bit binary string
        
    Raises:
        ValueError: If inputs are not 32 bits long
        
    Example:
        >>> xor32('0101'*8, '0011'*8)
        '01100110011001100110011001100110'
    """
    if len(x) != 32 or len(y) != 32:
        raise ValueError("Inputs must be 32 bits long")
    u = int(x, 2) ^ int(y, 2)
    return f'{u:032b}'

def xor128(x: str, y: str) -> str:
    """
    128-bit XOR operation between two binary strings.
    
    Args:
        x: First 128-bit binary string
        y: Second 128-bit binary string
        
    Returns:
        XOR result as 128-bit binary string
        
    Raises:
        ValueError: If inputs are not 128 bits long
        
    Example:
        >>> xor128('01'*64, '10'*64)
        '11'*64
    """
    if len(x) != 128 or len(y) != 128:
        raise ValueError("Inputs must be 128 bits long")
    u = int(x, 2) ^ int(y, 2)
    return f'{u:0128b}'

def Yand32(x: str, y: str) -> str:
    """
    32-bit AND operation between two binary strings.
    
    Args:
        x: First 32-bit binary string
        y: Second 32-bit binary string
        
    Returns:
        AND result as 32-bit binary string
        
    Raises:
        ValueError: If inputs are not 32 bits long
    """
    if len(x) != 32 or len(y) != 32:
        raise ValueError("Inputs must be 32 bits long")
    u = int(x, 2) & int(y, 2)
    return f'{u:032b}'

def pad_data(data: bytes, block_size: int = 128) -> bytes:
    """
    PKCS#7 padding for data to match block size.
    
    Args:
        data: Input data to pad
        block_size: Block size in bits
        
    Returns:
        Padded data
    """
    pad_len = block_size // 8 - (len(data) % (block_size // 8))
    return data + bytes([pad_len] * pad_len)

def unpad_data(data: bytes) -> bytes:
    """
    Remove PKCS#7 padding from data.
    
    Args:
        data: Padded data
        
    Returns:
        Unpadded data
        
    Raises:
        ValueError: If padding is invalid
    """
    pad_len = data[-1]
    if pad_len > len(data) or not all(b == pad_len for b in data[-pad_len:]):
        raise ValueError("Invalid padding")
    return data[:-pad_len]

def binary_to_bytes(binary_str: str) -> bytes:
    """
    Convert a binary string to bytes.
    
    Args:
        binary_str: Binary string (length must be multiple of 8)
        
    Returns:
        Corresponding bytes object
        
    Raises:
        ValueError: If input length isn't multiple of 8
    """
    if len(binary_str) % 8 != 0:
        raise ValueError("Binary string length must be multiple of 8")
    return bytes(int(binary_str[i:i+8], 2) for i in range(0, len(binary_str), 8))

def bytes_to_binary(data: bytes) -> str:
    """
    Convert bytes to a binary string.
    
    Args:
        data: Bytes to convert
        
    Returns:
        Binary string representation
    """
    return ''.join(f'{byte:08b}' for byte in data)

def generate_secure_key(key_length: int = 128) -> str:
    """
    Generate a cryptographically secure random key.
    
    Args:
        key_length: Length of key in bits
        
    Returns:
        Binary string representation of the key
    """
    # Use system random for cryptographic security
    sys_random = random.SystemRandom()
    return ''.join(str(sys_random.randint(0, 1)) for _ in range(key_length))

def validate_key(key: str) -> bool:
    """
    Validate that a key is a proper binary string of correct length.
    
    Args:
        key: Key to validate
        
    Returns:
        True if valid, False otherwise
    """
    if len(key) != 128:
        return False
    return all(c in ('0', '1') for c in key)

# ==================== CORE CIPHER FUNCTIONS ====================

def F(x: str, y: str) -> str:
    """
    Round function F used in the cipher.
    
    Args:
        x: 32-bit binary string
        y: 32-bit binary string (round key)
        
    Returns:
        32-bit result of the round function
        
    Raises:
        ValueError: If inputs are not 32 bits long
    """
    if len(x) != 32 or len(y) != 32:
        raise ValueError("Inputs must be 32 bits long")
    
    # XOR input with round key
    u = xor32(x, y)
    
    # Split into 8-bit chunks
    u1, u2, u3, u4 = u[0:8], u[8:16], u[16:24], u[24:32]
    
    # S-box substitution (custom S-box)
    S = {
        '0': 5, '1': 10, '2': 15, '3': 4, '4': 9, '5': 14, '6': 11, '7': 8,
        '8': 2, '9': 7, '10': 12, '11': 13, '12': 3, '13': 6, '14': 1, '15': 0
    }
    
    # Apply S-box to each nibble (4 bits)
    def apply_sbox(byte: str) -> str:
        return (format(S[str(int(byte[0:4], 2))], '04b') + 
                format(S[str(int(byte[4:8], 2))], '04b'))
    
    v1 = apply_sbox(u1)
    v2 = apply_sbox(u2)
    v3 = apply_sbox(u3)
    v4 = apply_sbox(u4)
    
    v = v1 + v2 + v3 + v4
    
    # Apply diffusion operations
    z = xor32(xor32(v, demo(v, 3)), demo(v, 10))
    return z

def ExpandKey(k: str) -> List[str]:
    """
    Key expansion algorithm to generate round keys.
    
    Args:
        k: 128-bit master key as binary string
        
    Returns:
        List of 96 round keys (each 32 bits)
        
    Raises:
        ValueError: If key is not 128 bits long
    """
    if len(k) != 128:
        raise ValueError("Key must be 128 bits long")
    
    kls = []
    # Initial 4 words from the key
    kls.append(k[0:32])
    kls.append(k[32:64])
    kls.append(k[64:96])
    kls.append(k[96:128])
    
    # Generate remaining 92 round keys
    for i in range(92):
        # Complex key expansion function
        term1 = xor32(kls[i], 32*'1')
        term2 = xor32(kls[i+1], 32*'1')
        term3 = kls[i+2]
        term4 = xor32(kls[i+3], format(i, '032b'))
        
        k0 = Yand32(Yand32(xor32(term1, term2), term3), term4)
        k0 = xor32(xor32(k0, demo(k0, 13)), demo(k0, 22))
        kls.append(k0)
    
    return kls

# ==================== ENCRYPTION/DECRYPTION ====================

def Enc(P: str, k: str) -> str:
    """
    Encrypt a 128-bit block of plaintext.
    
    Args:
        P: 128-bit plaintext block as binary string
        k: 128-bit master key as binary string
        
    Returns:
        128-bit ciphertext block as binary string
        
    Raises:
        ValueError: If inputs are not correct size
    """
    if len(P) != 128 or len(k) != 128:
        raise ValueError("Plaintext and key must be 128 bits long")
    
    kls = ExpandKey(k)
    als, bls, cls, dls = [], [], [], []
    
    # Initial state
    als.append(P[0:32])
    bls.append(P[32:64])
    cls.append(P[64:96])
    dls.append(P[96:128])
    
    # Main rounds (47 rounds)
    for i in range(47):
        als.append(xor32(F(als[i], kls[2*i]), bls[i]))
        cls.append(xor32(als[i+1], dls[i]))
        dls.append(xor32(F(dls[i], kls[2*i+1]), cls[i]))
        bls.append(xor32(dls[i+1], als[i]))
    
    # Final round
    bls.append(xor32(F(als[47], kls[94]), bls[47]))
    dls.append(xor32(bls[48], dls[47]))
    cls.append(xor32(F(dls[47], kls[95]), cls[47]))
    als.append(xor32(cls[48], als[47]))
    
    # Concatenate final state
    C = als[48] + bls[48] + cls[48] + dls[48]
    return C

def Dec(C: str, k: str) -> str:
    """
    Decrypt a 128-bit block of ciphertext.
    
    Args:
        C: 128-bit ciphertext block as binary string
        k: 128-bit master key as binary string
        
    Returns:
        128-bit plaintext block as binary string
        
    Raises:
        ValueError: If inputs are not correct size
    """
    if len(C) != 128 or len(k) != 128:
        raise ValueError("Ciphertext and key must be 128 bits long")
    
    kls = ExpandKey(k)
    als, bls, cls, dls = [], [], [], []
    
    # Initial state
    als.append(C[0:32])
    bls.append(C[32:64])
    cls.append(C[64:96])
    dls.append(C[96:128])
    
    # First inverse round
    als.append(xor32(als[0], cls[0]))
    bls.append(xor32(F(als[1], kls[94]), bls[0]))
    dls.append(xor32(bls[0], dls[0]))
    cls.append(xor32(F(dls[1], kls[95]), cls[0]))
    
    # Main inverse rounds (46 rounds)
    for i in range(1, 48):
        als.append(xor32(bls[i], dls[i]))
        bls.append(xor32(F(als[i+1], kls[2*(48-i)-2]), als[i]))
        dls.append(xor32(als[i], cls[i]))
        cls.append(xor32(F(dls[i+1], kls[2*(48-i)-1]), dls[i]))
    
    # Concatenate final state
    P = als[48] + bls[48] + cls[48] + dls[48]
    return P

# ==================== OPERATION MODES ====================

def ECB_encrypt(plaintext: bytes, key: str) -> bytes:
    """
    Encrypt data using Electronic Codebook (ECB) mode.
    
    Args:
        plaintext: Data to encrypt
        key: 128-bit master key as binary string
        
    Returns:
        Encrypted data
        
    Raises:
        ValueError: If key is invalid
    """
    if not validate_key(key):
        raise ValueError("Invalid key")
    
    # Pad data if needed
    padded_data = pad_data(plaintext)
    binary_data = bytes_to_binary(padded_data)
    
    # Process each 128-bit block
    ciphertext = []
    for i in range(0, len(binary_data), 128):
        block = binary_data[i:i+128]
        encrypted_block = Enc(block, key)
        ciphertext.append(encrypted_block)
    
    return binary_to_bytes(''.join(ciphertext))

def ECB_decrypt(ciphertext: bytes, key: str) -> bytes:
    """
    Decrypt data using Electronic Codebook (ECB) mode.
    
    Args:
        ciphertext: Data to decrypt
        key: 128-bit master key as binary string
        
    Returns:
        Decrypted data
        
    Raises:
        ValueError: If key is invalid or ciphertext is malformed
    """
    if not validate_key(key):
        raise ValueError("Invalid key")
    
    binary_data = bytes_to_binary(ciphertext)
    if len(binary_data) % 128 != 0:
        raise ValueError("Ciphertext length must be multiple of 128 bits")
    
    # Process each 128-bit block
    plaintext = []
    for i in range(0, len(binary_data), 128):
        block = binary_data[i:i+128]
        decrypted_block = Dec(block, key)
        plaintext.append(decrypted_block)
    
    # Remove padding
    return unpad_data(binary_to_bytes(''.join(plaintext)))

def CBC_encrypt(plaintext: bytes, key: str, iv: str = None) -> bytes:
    """
    Encrypt data using Cipher Block Chaining (CBC) mode.
    
    Args:
        plaintext: Data to encrypt
        key: 128-bit master key as binary string
        iv: 128-bit initialization vector (random if None)
        
    Returns:
        Tuple of (ciphertext, iv_used)
        
    Raises:
        ValueError: If key is invalid or iv is malformed
    """
    if not validate_key(key):
        raise ValueError("Invalid key")
    
    if iv is None:
        iv = generate_secure_key()
    elif len(iv) != 128 or not all(c in ('0', '1') for c in iv):
        raise ValueError("IV must be 128-bit binary string")
    
    # Pad data if needed
    padded_data = pad_data(plaintext)
    binary_data = bytes_to_binary(padded_data)
    
    # Process blocks with CBC mode
    ciphertext = []
    prev_block = iv
    
    for i in range(0, len(binary_data), 128):
        block = binary_data[i:i+128]
        xored_block = xor128(block, prev_block)
        encrypted_block = Enc(xored_block, key)
        ciphertext.append(encrypted_block)
        prev_block = encrypted_block
    
    return binary_to_bytes(''.join(ciphertext)), iv

def CBC_decrypt(ciphertext: bytes, key: str, iv: str) -> bytes:
    """
    Decrypt data using Cipher Block Chaining (CBC) mode.
    
    Args:
        ciphertext: Data to decrypt
        key: 128-bit master key as binary string
        iv: 128-bit initialization vector used in encryption
        
    Returns:
        Decrypted data
        
    Raises:
        ValueError: If key/iv is invalid or ciphertext is malformed
    """
    if not validate_key(key):
        raise ValueError("Invalid key")
    if len(iv) != 128 or not all(c in ('0', '1') for c in iv):
        raise ValueError("IV must be 128-bit binary string")
    
    binary_data = bytes_to_binary(ciphertext)
    if len(binary_data) % 128 != 0:
        raise ValueError("Ciphertext length must be multiple of 128 bits")
    
    # Process blocks with CBC mode
    plaintext = []
    prev_block = iv
    
    for i in range(0, len(binary_data), 128):
        block = binary_data[i:i+128]
        decrypted_block = Dec(block, key)
        xored_block = xor128(decrypted_block, prev_block)
        plaintext.append(xored_block)
        prev_block = block
    
    # Remove padding
    return unpad_data(binary_to_bytes(''.join(plaintext)))

# ==================== IMAGE ENCRYPTION/DECRYPTION ====================

def encrypt_image(input_path: str, output_path: str, key: str = None, mode: str = 'ECB') -> Tuple[str, str]:
    """
    Encrypt an image file using specified mode.
    
    Args:
        input_path: Path to input image
        output_path: Path to save encrypted image
        key: Encryption key (random if None)
        mode: 'ECB' or 'CBC'
        
    Returns:
        Tuple of (key_used, iv_used) (iv is None for ECB)
        
    Raises:
        ValueError: For invalid inputs
        IOError: For file operations
    """
    if mode not in ('ECB', 'CBC'):
        raise ValueError("Mode must be 'ECB' or 'CBC'")
    
    try:
        img = Image.open(input_path)
        P_arr = np.array(img)
    except Exception as e:
        raise IOError(f"Error opening image: {str(e)}")
    
    # Convert image to binary string
    P = ''
    for i in range(P_arr.shape[0]):
        for j in range(P_arr.shape[1]):
            for k in range(P_arr.shape[2]):
                P += format(int(P_arr[i][j][k]), '08b')
    
    # Generate key if not provided
    if key is None:
        key = generate_secure_key()
    elif not validate_key(key):
        raise ValueError("Invalid key format")
    
    # Encrypt based on mode
    if mode == 'ECB':
        C = ''
        for i in range(0, len(P), 128):
            C += Enc(P[i:i+128], key)
        iv = None
    else:  # CBC
        iv = generate_secure_key()
        C = Enc(xor128(P[0:128], iv), key)
        for i in range(1, len(P) // 128):
            C += Enc(xor128(P[128*i:128*(i+1)], C[128*(i-1):128*i]), key)
    
    # Convert ciphertext back to image
    C_arr = np.array(img)
    byte_index = 0
    for i in range(C_arr.shape[0]):
        for j in range(C_arr.shape[1]):
            for k in range(C_arr.shape[2]):
                if byte_index + 8 <= len(C):
                    C_arr[i][j][k] = int(C[byte_index:byte_index+8], 2)
                    byte_index += 8
                else:
                    C_arr[i][j][k] = 0  # Pad with zeros if needed
    
    # Save encrypted image
    try:
        encrypted_img = Image.fromarray(C_arr.astype('uint8'))
        encrypted_img.save(output_path)
    except Exception as e:
        raise IOError(f"Error saving encrypted image: {str(e)}")
    
    return key, iv

def decrypt_image(input_path: str, output_path: str, key: str, iv: str = None, mode: str = 'ECB') -> None:
    """
    Decrypt an image file using specified mode.
    
    Args:
        input_path: Path to encrypted image
        output_path: Path to save decrypted image
        key: Encryption key used
        iv: Initialization vector (required for CBC)
        mode: 'ECB' or 'CBC'
        
    Raises:
        ValueError: For invalid inputs
        IOError: For file operations
    """
    if mode not in ('ECB', 'CBC'):
        raise ValueError("Mode must be 'ECB' or 'CBC'")
    if not validate_key(key):
        raise ValueError("Invalid key format")
    if mode == 'CBC' and (iv is None or len(iv) != 128 or not all(c in ('0', '1') for c in iv)):
        raise ValueError("IV must be provided for CBC mode and be 128-bit binary string")
    
    try:
        img = Image.open(input_path)
        C_arr = np.array(img)
    except Exception as e:
        raise IOError(f"Error opening image: {str(e)}")
    
    # Convert image to binary string
    C = ''
    for i in range(C_arr.shape[0]):
        for j in range(C_arr.shape[1]):
            for k in range(C_arr.shape[2]):
                C += format(int(C_arr[i][j][k]), '08b')
    
    # Decrypt based on mode
    if mode == 'ECB':
        P = ''
        for i in range(0, len(C), 128):
            P += Dec(C[i:i+128], key)
    else:  # CBC
        P = xor128(Dec(C[0:128], key), iv)
        for i in range(1, len(C) // 128):
            P += xor128(Dec(C[128*i:128*(i+1)], key), C[128*(i-1):128*i])
    
    # Convert plaintext back to image
    P_arr = np.array(img)
    byte_index = 0
    for i in range(P_arr.shape[0]):
        for j in range(P_arr.shape[1]):
            for k in range(P_arr.shape[2]):
                if byte_index + 8 <= len(P):
                    P_arr[i][j][k] = int(P[byte_index:byte_index+8], 2)
                    byte_index += 8
                else:
                    P_arr[i][j][k] = 0  # Pad with zeros if needed
    
    # Save decrypted image
    try:
        decrypted_img = Image.fromarray(P_arr.astype('uint8'))
        decrypted_img.save(output_path)
    except Exception as e:
        raise IOError(f"Error saving decrypted image: {str(e)}")

# ==================== TESTING AND EXAMPLES ====================

def run_self_tests():
    """Run comprehensive self-tests to verify all functionality."""
    print("Running self-tests...")
    
    # Test utility functions
    assert demo('101100', 2) == '110010'
    assert xor32('0101'*8, '0011'*8) == '01100110011001100110011001100110'
    assert Yand32('0101'*8, '0011'*8) == '00010001000100010001000100010001'
    
    # Test padding
    test_data = b'test'
    padded = pad_data(test_data)
    assert len(padded) == 16  # 128 bits
    assert unpad_data(padded) == test_data
    
    # Test key generation and validation
    test_key = generate_secure_key()
    assert validate_key(test_key)
    assert not validate_key(test_key[:-1] + '2')  # Invalid character
    
    # Test core cipher functions
    test_key = '0'*128
    expanded = ExpandKey(test_key)
    assert len(expanded) == 96  # 96 round keys
    
    # Test encryption/decryption round trip
    test_plaintext = '0'*128
    ciphertext = Enc(test_plaintext, test_key)
    decrypted = Dec(ciphertext, test_key)
    assert decrypted == test_plaintext
    
    # Test ECB mode
    test_data = b'This is a test message for ECB mode!'
    encrypted = ECB_encrypt(test_data, test_key)
    decrypted = ECB_decrypt(encrypted, test_key)
    assert decrypted == test_data
    
    # Test CBC mode
    iv = generate_secure_key()
    encrypted, _ = CBC_encrypt(test_data, test_key, iv)
    decrypted = CBC_decrypt(encrypted, test_key, iv)
    assert decrypted == test_data
    
    print("All self-tests passed successfully!")

def example_usage():
    """Demonstrate example usage of the cryptographic system."""
    print("\nExample Usage:")
    
    # Generate a secure key
    key = generate_secure_key()
    print(f"Generated key: {key[:32]}... (length: {len(key)} bits)")
    
    # Sample data to encrypt
    message = b"Hello, this is a secret message!"
    print(f"\nOriginal message: {message.decode()}")
    
    # ECB mode example
    print("\nECB Mode:")
    ecb_encrypted = ECB_encrypt(message, key)
    print(f"Encrypted (hex): {ecb_encrypted.hex()}")
    ecb_decrypted = ECB_decrypt(ecb_encrypted, key)
    print(f"Decrypted: {ecb_decrypted.decode()}")
    
    # CBC mode example
    print("\nCBC Mode:")
    cbc_encrypted, iv = CBC_encrypt(message, key)
    print(f"IV used: {iv[:32]}...")
    print(f"Encrypted (hex): {cbc_encrypted.hex()}")
    cbc_decrypted = CBC_decrypt(cbc_encrypted, key, iv)
    print(f"Decrypted: {cbc_decrypted.decode()}")
    
    # Image encryption example
    print("\nImage Encryption:")
    try:
        # This would actually need an image file to work
        # key_used, iv_used = encrypt_image('input.bmp', 'encrypted.bmp', mode='CBC')
        # decrypt_image('encrypted.bmp', 'decrypted.bmp', key_used, iv_used, mode='CBC')
        print("Image encryption/decryption functions available (commented out for this demo)")
    except Exception as e:
        print(f"Image example error: {str(e)}")

if __name__ == "__main__":
    run_self_tests()
    example_usage()